from PIL import Image
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import classification.cnn.cnn as model
from classification.dataPreparation.load_data import get_files
import classification.dataPreparation.save_data as data
from classification.config.Params import Params


params = Params()

# =======================================================================
# 获取一张图片
def get_one_image(test_image, label):
    # 输入参数：train,训练图片的路径
    # 返回参数：image，从训练图片中随机抽取一张图片
    n = len(test_image)
    ind = np.random.randint(0, n)
    img_dir = test_image[ind]  # 随机选择测试的图片
    index = int(img_dir.split("_")[3].split(".")[0])
    label = params.classes.get(index)
    img = Image.open(img_dir)
    image = np.array(img)
    return image, label, index


# --------------------------------------------------------------------
# 测试图片
def evaluate_one_image(image_array, shape):
    width = shape[0]
    length = shape[1]
    with tf.Graph().as_default():
        image = tf.cast(image_array, tf.float32)
        image = tf.image.per_image_standardization(image)
        image = tf.reshape(image, [1, width, length, 3])

        logit = model.inference(image, 1, params.n_classes)

        logit = tf.nn.softmax(logit)

        x = tf.placeholder(tf.float32, shape=[width, length, 3])

        logs_train_dir = params.logs_train_path

        saver = tf.train.Saver()

        with tf.Session() as sess:

            print("Reading checkpoints...")
            ckpt = tf.train.get_checkpoint_state(logs_train_dir)
            if ckpt and ckpt.model_checkpoint_path:
                global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
                saver.restore(sess, ckpt.model_checkpoint_path)
                print('Loading success, global_step is %s' % global_step)
            else:
                print('No checkpoint file found')

            prediction = sess.run(logit, feed_dict={x: image_array})
            max_index = np.argmax(prediction)
            if max_index == 0:
                print('This is a daisy with possibility %.6f' % prediction[:, 0])
            elif max_index == 1:
                print('This is a dandelion with possibility %.6f' % prediction[:, 1])
            elif max_index == 2:
                print('This is a rose with possibility %.6f' % prediction[:, 2])
            elif max_index == 3:
                print('This is a sunflower with possibility %.6f' % prediction[:, 3])
            else:
                print('This is a tulip with possibility %.6f' % prediction[:, 4])

    return max_index

if __name__ == '__main__':
    train_dir = params.samples_path
    train, train_label, val, val_label = get_files(train_dir)
    right = 0
    for i in range(100):
        img, label, index = get_one_image(val, val_label)
        evaluate_one_image(img, params.shape)
        if evaluate_one_image(img, params.shape) == index:
            right = right + 1
        print("this flower is :" + label)
    print("right number:%d" % right)
